from ryu.base import app_manager
from ryu.controller import ofp_event
from ryu.controller.handler import CONFIG_DISPATCHER, MAIN_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.ofproto import ofproto_v1_3
from ryu.lib.packet import packet
from ryu.lib.packet import arp
from ryu.lib.packet import ethernet
from ryu.lib.packet import ether_types

from ryu.lib.packet import in_proto
from ryu.lib.packet import ipv4
from ryu.lib.packet import icmp
from ryu.lib.packet import tcp
from ryu.lib.packet import udp

FLOW_SERIAL_NO = 0

def get_flow_number():
    global FLOW_SERIAL_NO
    FLOW_SERIAL_NO = FLOW_SERIAL_NO + 1
    return FLOW_SERIAL_NO


class SimpleSwitch13(app_manager.RyuApp):
    OFP_VERSIONS = [ofproto_v1_3.OFP_VERSION]

    def __init__(self, *args, **kwargs):
        super(SimpleSwitch13, self).__init__(*args, **kwargs)
        self.mac_to_port = {}
        self.arp_ip_to_port = {} # Stores {dpid: {in_port: [ip1, ip2,...]}}
        self.mitigation = True # Enable mitigation by default

    @set_ev_cls(ofp_event.EventOFPSwitchFeatures, CONFIG_DISPATCHER)
    def switch_features_handler(self, ev):
        datapath = ev.msg.datapath
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser
        dpid = datapath.id

        self.mac_to_port.setdefault(dpid, {})
        self.arp_ip_to_port.setdefault(dpid, {})

        flow_serial_no = get_flow_number()
        match = parser.OFPMatch()
        actions = [parser.OFPActionOutput(ofproto.OFPP_CONTROLLER,
                                          ofproto.OFPCML_NO_BUFFER)]
        self.add_flow(datapath, 0, match, actions, flow_serial_no)
        self.logger.info("Switch %s connected. Installed table-miss flow.", dpid)

    def add_flow(self, datapath, priority, match, actions, serial_no, buffer_id=None, idle=0, hard=0):
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser

        inst = [parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS,
                                             actions)]
        if buffer_id:
            mod = parser.OFPFlowMod(datapath=datapath, cookie=serial_no, buffer_id=buffer_id,
                                    idle_timeout=idle, hard_timeout=hard,
                                    priority=priority, match=match,
                                    instructions=inst)
        else:
            mod = parser.OFPFlowMod(datapath=datapath, cookie=serial_no, priority=priority,
                                    idle_timeout=idle, hard_timeout=hard,
                                    match=match, instructions=inst)
        datapath.send_msg(mod)

    def add_ip_drop_flow(self, datapath, in_port, src_ip, hard_timeout=60):
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser
        match = parser.OFPMatch(in_port=in_port,
                                eth_type=ether_types.ETH_TYPE_IP,
                                ipv4_src=src_ip)
        actions = []
        flow_serial_no = get_flow_number()
        self.logger.info("MITIGATION: Installing drop flow for IP %s on dpid %s port %s for %s seconds.",
                         src_ip, datapath.id, in_port, hard_timeout)
        self.add_flow(datapath, 100, match, actions, flow_serial_no, hard=hard_timeout)

    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def _packet_in_handler(self, ev):
        if ev.msg.msg_len < ev.msg.total_len:
            self.logger.debug("packet truncated: only %s of %s bytes",
                              ev.msg.msg_len, ev.msg.total_len)
        msg = ev.msg
        datapath = msg.datapath
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser
        in_port = msg.match['in_port']
        dpid = datapath.id

        pkt = packet.Packet(msg.data)
        eth = pkt.get_protocols(ethernet.ethernet)[0]

        if eth.ethertype == ether_types.ETH_TYPE_LLDP:
            return
        
        dst = eth.dst
        src = eth.src

        self.arp_ip_to_port[dpid].setdefault(in_port, [])

        self.mac_to_port[dpid][src] = in_port

        if dst in self.mac_to_port[dpid]:
            out_port = self.mac_to_port[dpid][dst]
        else:
            out_port = ofproto.OFPP_FLOOD

        actions = [parser.OFPActionOutput(out_port)]

        if eth.ethertype == ether_types.ETH_TYPE_ARP:
            arp_pkt = pkt.get_protocol(arp.arp)
            if arp_pkt.opcode == arp.ARP_REQUEST or arp_pkt.opcode == arp.ARP_REPLY:
                if arp_pkt.src_ip not in self.arp_ip_to_port[dpid][in_port]:
                    self.arp_ip_to_port[dpid][in_port].append(arp_pkt.src_ip)
                    self.logger.info("ARP_LEARN: IP %s (MAC %s) seen on dpid %s port %s via ARP.",
                                     arp_pkt.src_ip, arp_pkt.src_mac, dpid, in_port)

        if out_port != ofproto.OFPP_FLOOD:
            match_flow = None

            if eth.ethertype == ether_types.ETH_TYPE_IP:
                ip = pkt.get_protocol(ipv4.ipv4)
                srcip = ip.src
                dstip = ip.dst
                protocol = ip.proto

                if self.mitigation:
                    if srcip not in self.arp_ip_to_port[dpid][in_port]:
                        self.logger.warning(
                            "MITIGATION: Suspicious IP packet from %s (MAC %s) on dpid %s port %s. "
                            "This IP was not previously announced via ARP on this port.",
                            srcip, src, dpid, in_port)
                        self.add_ip_drop_flow(datapath, in_port, srcip, hard_timeout=60)
                        return

                match_flow = parser.OFPMatch(eth_type=ether_types.ETH_TYPE_IP,
                                             ipv4_src=srcip, ipv4_dst=dstip,
                                             ip_proto=protocol)
                
                if protocol == in_proto.IPPROTO_ICMP:
                    icmp_pkt = pkt.get_protocol(icmp.icmp)
                    match_flow = parser.OFPMatch(in_port=in_port, eth_dst=dst, eth_src=src,
                                                 eth_type=ether_types.ETH_TYPE_IP,
                                                 ipv4_src=srcip, ipv4_dst=dstip,
                                                 ip_proto=protocol,
                                                 icmpv4_code=icmp_pkt.code,
                                                 icmpv4_type=icmp_pkt.type)
                elif protocol == in_proto.IPPROTO_TCP:
                    tcp_pkt = pkt.get_protocol(tcp.tcp)
                    match_flow = parser.OFPMatch(in_port=in_port, eth_dst=dst, eth_src=src,
                                                 eth_type=ether_types.ETH_TYPE_IP,
                                                 ipv4_src=srcip, ipv4_dst=dstip,
                                                 ip_proto=protocol,
                                                 tcp_src=tcp_pkt.src_port, tcp_dst=tcp_pkt.dst_port)
                elif protocol == in_proto.IPPROTO_UDP:
                    udp_pkt = pkt.get_protocol(udp.udp)
                    match_flow = parser.OFPMatch(in_port=in_port, eth_dst=dst, eth_src=src,
                                                 eth_type=ether_types.ETH_TYPE_IP,
                                                 ipv4_src=srcip, ipv4_dst=dstip,
                                                 ip_proto=protocol,
                                                 udp_src=udp_pkt.src_port, udp_dst=udp_pkt.dst_port)
            
            elif eth.ethertype == ether_types.ETH_TYPE_ARP:
                 match_flow = parser.OFPMatch(in_port=in_port, eth_dst=dst, eth_src=src,
                                             eth_type=ether_types.ETH_TYPE_ARP)

            if match_flow:
                flow_serial_no = get_flow_number()
                if msg.buffer_id != ofproto.OFP_NO_BUFFER:
                    self.add_flow(datapath, 1, match_flow, actions, flow_serial_no, msg.buffer_id, idle=20, hard=100)
                else:
                    self.add_flow(datapath, 1, match_flow, actions, flow_serial_no, idle=20, hard=100)
                if msg.buffer_id != ofproto.OFP_NO_BUFFER:
                    return

        data = None
        if msg.buffer_id == ofproto.OFP_NO_BUFFER:
            data = msg.data

        out = parser.OFPPacketOut(datapath=datapath, buffer_id=msg.buffer_id,
                                  in_port=in_port, actions=actions, data=data)
        datapath.send_msg(out)